import warnings
import tushare as ts
import pandas as pd
import numpy as np
import xgboost as xgb
from sklearn.model_selection import train_test_split, GridSearchCV
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
import os

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
warnings.filterwarnings("ignore")


def initialize_tushare(token):
    """初始化 Tushare API."""
    return ts.pro_api(token)


def clean_data(df):
    """清洗数据，处理异常值和缺失值."""
    df = df.copy()
    df.fillna(method='ffill', inplace=True)  # 前向填充缺失值
    df['pct_change'] = df['close'].pct_change()
    df = df[~(df['pct_change'].abs() > 0.1)]  # 移除日价格变化超过10%的记录
    df.drop(columns=['pct_change'], inplace=True)
    return df


def get_stock_data(pro, stocks, start_date, end_date):
    """获取多个股票的数据并返回字典格式的数据."""
    data_dict = {}
    for stock in stocks:
        df = pro.daily(ts_code=stock, start_date=start_date, end_date=end_date)
        df = df[['trade_date', 'open', 'high', 'low', 'close', 'vol']]
        df['trade_date'] = pd.to_datetime(df['trade_date'])
        df.set_index('trade_date', inplace=True)
        df.sort_index(inplace=True)
        df = clean_data(df)  # 数据清洗
        data_dict[stock] = df
    return data_dict


def traditional_strategy(df):
    """执行传统策略，生成买入和卖出信号."""
    df['ma5'] = df['close'].rolling(window=5).mean()
    df['ma20'] = df['close'].rolling(window=20).mean()
    df['20d_high'] = df['high'].rolling(window=20).max()
    df['20d_low'] = df['low'].rolling(window=20).min()
    df['vol_ma5'] = df['vol'].rolling(window=5).mean()

    df['Signal'] = 0
    df.loc[(df['close'] > df['20d_high'].shift(1)) &
           (df['ma5'] > df['ma20']) &
           (df['vol'] > df['vol_ma5'] * 0.8), 'Signal'] = 1  # 买入信号
    df.loc[(df['close'] < df['20d_low'].shift(1)) &
           (df['ma5'] < df['ma20']) &
           (df['vol'] < df['vol_ma5'] * 1.2), 'Signal'] = -1  # 卖出信号
    return df


def plot_traditional_signals(df, stock):
    """绘制传统策略交易信号示意图（图 3.1）."""
    plt.figure(figsize=(10, 5))
    plt.plot(df.index, df['close'], label="价格", color="black", linewidth=1.5)
    buy_signals = df[df['Signal'] == 1]
    sell_signals = df[df['Signal'] == -1]
    plt.scatter(buy_signals.index, buy_signals['close'], color="green", label="买入信号", marker="^", s=100)
    plt.scatter(sell_signals.index, sell_signals['close'], color="red", label="卖出信号", marker="v", s=100)
    plt.xlabel("交易日")
    plt.ylabel("价格 (元)")
    plt.title(f"图 3.1 传统策略交易信号示意图 - {stock}")
    plt.legend()
    plt.grid(True, linestyle="--", alpha=0.7)
    plt.gcf().autofmt_xdate()  # 自动格式化日期标签
    plt.tight_layout()
    plt.savefig(f"traditional_signals_{stock}.png", dpi=300)
    plt.close()


def plot_ml_signals(df, stock):
    """绘制机器学习策略信号生成示意图（图 3.2）."""
    plt.figure(figsize=(10, 5))
    plt.plot(df.index, df['close'], label="价格", color="black", linewidth=1.5)
    buy_signals = df[df['Signal_ML'] == 1]
    sell_signals = df[df['Signal_ML'] == -1]
    plt.scatter(buy_signals.index, buy_signals['close'], color="green", label="买入信号", marker="^", s=100)
    plt.scatter(sell_signals.index, sell_signals['close'], color="red", label="卖出信号", marker="v", s=100)
    plt.xlabel("交易日")
    plt.ylabel("价格 (元)")
    plt.title(f"图 3.2 机器学习策略信号生成示意图 - {stock}")
    plt.legend()
    plt.grid(True, linestyle="--", alpha=0.7)
    plt.text(df.index[-20], df['close'].min(), "5 天交易冷却期", fontsize=10, bbox=dict(facecolor='yellow', alpha=0.2))
    plt.gcf().autofmt_xdate()
    plt.tight_layout()
    plt.savefig(f"ml_signals_{stock}.png", dpi=300)
    plt.close()


def add_features(df):
    """为数据框添加特征并生成目标变量."""

    def calculate_rsi(data, periods=14):
        delta = data['close'].diff()
        gain = (delta.where(delta > 0, 0)).rolling(window=periods).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(window=periods).mean()
        rs = gain / loss
        return 100 - (100 / (1 + rs))

    df['RSI'] = calculate_rsi(df)
    df['ema12'] = df['close'].ewm(span=12, adjust=False).mean()
    df['ema26'] = df['close'].ewm(span=26, adjust=False).mean()
    df['macd'] = df['ema12'] - df['ema26']
    df['signal_line'] = df['macd'].ewm(span=9, adjust=False).mean()

    df['ma20'] = df['close'].rolling(window=20).mean()
    df['std20'] = df['close'].rolling(window=20).std()
    df['bb_upper'] = df['ma20'] + 2 * df['std20']
    df['bb_lower'] = df['ma20'] - 2 * df['std20']

    df['pct_change_1'] = df['close'].pct_change(1)
    df['pct_change_5'] = df['close'].pct_change(5)
    df['pct_change_10'] = df['close'].pct_change(10)

    df['high_low'] = df['high'] - df['low']
    df['high_close'] = np.abs(df['high'] - df['close'].shift(1))
    df['low_close'] = np.abs(df['low'] - df['close'].shift(1))
    df['tr'] = df[['high_low', 'high_close', 'low_close']].max(axis=1)
    df['ATR'] = df['tr'].rolling(window=14).mean()

    df['direction'] = np.where(df['close'] > df['close'].shift(1), 1, -1)
    df['OBV'] = (df['direction'] * df['vol']).cumsum()

    df['Target'] = (df['close'].shift(-1) > df['close']).astype(int)

    df.dropna(inplace=True)
    return df


def ml_strategy_with_tuning(df):
    """机器学习策略以生成交易信号."""
    df = add_features(df)

    features = df[['close', 'vol', 'RSI', 'macd', 'signal_line', 'bb_upper',
                   'bb_lower', 'pct_change_1', 'pct_change_5', 'pct_change_10',
                   'ATR', 'OBV']]
    target = df['Target']

    features = features.loc[target.notna()]
    target = target.dropna()

    X_train, X_test, y_train, y_test = train_test_split(features, target, test_size=0.2, random_state=42)

    model = xgb.XGBClassifier(use_label_encoder=False, eval_metric='logloss', random_state=42)

    param_grid = {
        'max_depth': [3, 4, 5],
        'learning_rate': [0.01, 0.1],
        'n_estimators': [100, 200],
        'subsample': [0.6, 0.8],
        'colsample_bytree': [0.6, 0.8]
    }

    grid_search = GridSearchCV(estimator=model, param_grid=param_grid,
                               scoring='accuracy', cv=3, verbose=1, n_jobs=-1)

    grid_search.fit(X_train, y_train)

    best_params = grid_search.best_params_
    print("最佳参数:", best_params)

    signals = grid_search.best_estimator_.predict_proba(features)[:, 1]
    df['Signal_ML'] = 0
    cooldown = 5  # 5天冷却期
    last_trade = -cooldown - 1
    for i in range(len(signals)):
        if i - last_trade <= cooldown:
            continue
        if signals[i] > 0.7:
            df.iloc[i, df.columns.get_loc('Signal_ML')] = 1
            last_trade = i
        elif signals[i] < 0.3:
            df.iloc[i, df.columns.get_loc('Signal_ML')] = -1
            last_trade = i

    return df, grid_search.best_estimator_


def backtest(df, strategy='traditional', initial_capital=1000, transaction_cost=0.002, slippage=0.001):
    """根据交易信号回测策略并计算投资组合价值."""
    position = 0
    cash = initial_capital
    portfolio_value = [initial_capital]
    signal_col = 'Signal' if strategy == 'traditional' else 'Signal_ML'
    entry_price = 0

    for i in range(len(df)):
        if i == 0:
            continue
        current_price = df['close'].iloc[i]
        adjusted_price = current_price * (1 + slippage if position == 0 else 1 - slippage)

        if df[signal_col].iloc[i] == 1 and cash > 0:
            cost = adjusted_price * (1 + transaction_cost)
            position = cash / cost
            cash = 0
            entry_price = current_price
        elif df[signal_col].iloc[i] == -1 and position > 0:
            proceeds = position * adjusted_price * (1 - transaction_cost)
            cash = proceeds
            position = 0
            entry_price = 0

        if position > 0:
            stop_loss = entry_price * (1 - 0.03)  # 3% 止损
            take_profit = entry_price * (1 + 0.05)  # 5% 止盈
            if current_price < stop_loss or current_price > take_profit:
                proceeds = position * adjusted_price * (1 - transaction_cost)
                cash = proceeds
                position = 0
                entry_price = 0

        value = cash + position * current_price if position > 0 else cash
        portfolio_value.append(value)

    return_rate = (portfolio_value[-1] - initial_capital) / initial_capital
    daily_returns = pd.Series(portfolio_value).pct_change().dropna()
    sharpe_ratio = daily_returns.mean() / daily_returns.std() * np.sqrt(252) if daily_returns.std() != 0 else 0
    return return_rate, sharpe_ratio, portfolio_value


def plot_data_analysis(data_dict):
    """绘制数据统计分析示例图（图 2.1）."""
    stats = []
    stocks = list(data_dict.keys())
    for stock, df in data_dict.items():
        avg_close = df['close'].mean()
        volatility = df['close'].std()
        stats.append([stock, avg_close, volatility])
    stats_df = pd.DataFrame(stats, columns=['股票代码', '平均收盘价 (元)', '波动率 (标准差)'])

    plt.figure(figsize=(8, 5))
    plt.bar(stats_df['股票代码'], stats_df['平均收盘价 (元)'], color='skyblue', label='平均收盘价')
    plt.plot(stats_df['股票代码'], stats_df['波动率 (标准差)'] * 100, color='red', marker='o', label='波动率 (x100)')
    plt.xlabel("股票代码")
    plt.ylabel("数值")
    plt.title("图 2.1 数据统计分析示例")
    plt.legend()
    plt.grid(True, linestyle="--", alpha=0.7)
    plt.xticks(rotation=45)
    plt.tight_layout()
    plt.savefig("data_analysis.png", dpi=300)
    plt.close()


def plot_portfolio_comparison(portfolio_values, data_dict):
    """为每只股票绘制传统策略和机器学习策略的投资组合价值对比图（图 5.1 至 图 5.8）."""
    if not os.path.exists('charts'):
        os.makedirs('charts')
    for stock, values in portfolio_values.items():
        dates = data_dict[stock].index
        plt.figure(figsize=(10, 6))
        plt.plot(dates, values['Traditional'], label='传统策略', color='blue', linewidth=2)
        plt.plot(dates, values['ML'], label='机器学习策略', color='orange', linewidth=2)
        plt.title(f"图 5.{list(portfolio_values.keys()).index(stock) + 1} {stock} 投资组合价值对比")
        plt.xlabel("交易日")
        plt.ylabel("投资组合价值 (元)")
        plt.legend()
        plt.grid(True, linestyle='--', alpha=0.7)
        plt.gcf().autofmt_xdate()  # 自动格式化日期标签
        plt.tight_layout()
        plt.savefig(f"portfolio_comparison_{stock}.png", dpi=300)
        plt.close()


def plot_specific_returns_comparison(results, stock='601988.SH'):
    """绘制特定股票的收益率对比图（图 4.1）."""
    trad_return = results[stock]['Traditional']['return'] * 100
    ml_return = results[stock]['ML']['return'] * 100

    plt.figure(figsize=(6, 4))
    plt.bar(['传统策略', '机器学习策略'], [trad_return, ml_return], color=['blue', 'orange'])
    plt.xlabel("策略")
    plt.ylabel("收益率 (%)")
    plt.title(f"图 4.1 {stock} 收益率对比")
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig(f"returns_comparison_{stock}.png", dpi=300)
    plt.close()


def plot_returns_comparison(results):
    """绘制所有股票的收益率对比图（图 5.9）."""
    if not os.path.exists('charts'):
        os.makedirs('charts')
    stocks_list = list(results.keys())
    trad_returns = [results[stock]['Traditional']['return'] * 100 for stock in stocks_list]
    ml_returns = [results[stock]['ML']['return'] * 100 for stock in stocks_list]

    x = np.arange(len(stocks_list))
    width = 0.35

    plt.figure(figsize=(12, 6))
    plt.bar(x - width / 2, trad_returns, width, label='传统策略', color='blue')
    plt.bar(x + width / 2, ml_returns, width, label='机器学习策略', color='orange')
    plt.xlabel("股票代码")
    plt.ylabel("收益率 (%)")
    plt.title("图 5.9 各股票收益率对比")
    plt.xticks(x, stocks_list, rotation=45)
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig("returns_comparison.png", dpi=300)
    plt.close()


def plot_sharpe_comparison(results):
    """绘制所有股票的夏普比率对比图（图 5.10）."""
    if not os.path.exists('charts'):
        os.makedirs('charts')
    stocks_list = list(results.keys())
    trad_sharpes = [results[stock]['Traditional']['sharpe'] for stock in stocks_list]
    ml_sharpes = [results[stock]['ML']['sharpe'] for stock in stocks_list]

    x = np.arange(len(stocks_list))
    width = 0.35

    plt.figure(figsize=(12, 6))
    plt.bar(x - width / 2, trad_sharpes, width, label='传统策略', color='blue')
    plt.bar(x + width / 2, ml_sharpes, width, label='机器学习策略', color='orange')
    plt.xlabel("股票代码")
    plt.ylabel("夏普比率")
    plt.title("图 5.10 各股票夏普比率对比")
    plt.xticks(x, stocks_list, rotation=45)
    plt.legend()
    plt.grid(True, linestyle='--', alpha=0.7)
    plt.tight_layout()
    plt.savefig("sharpe_comparison.png", dpi=300)
    plt.close()


def generate_portfolio_description(portfolio_values, results):
    """根据投资组合价值数据生成描述."""
    descriptions = {}
    for stock, values in portfolio_values.items():
        trad_values = values['Traditional']
        ml_values = values['ML']
        trad_return = results[stock]['Traditional']['return']
        ml_return = results[stock]['ML']['return']

        # 计算初期、中期、后期的价值
        n = len(trad_values)
        initial_trad = trad_values[n // 3]
        mid_trad = trad_values[2 * n // 3]
        final_trad = trad_values[-1]
        initial_ml = ml_values[n // 3]
        mid_ml = ml_values[2 * n // 3]
        final_ml = ml_values[-1]

        # 分析趋势并生成描述
        trad_trend = "稳定" if abs(
            final_trad - initial_trad) < 100 else "缓慢增长" if final_trad > initial_trad else "下降"
        ml_trend = "稳定" if abs(final_ml - initial_ml) < 100 else "显著增长" if final_ml > initial_ml else "下降"

        if trad_trend == "稳定":
            trad_desc = f"传统策略保持在 {int(final_trad)} 元左右，收益率 {trad_return:.2%}。"
        else:
            trad_desc = f"传统策略在初期{trad_trend}，中期达到约 {int(mid_trad)} 元，最终达到约 {int(final_trad)} 元，收益率 {trad_return:.2%}。"

        if ml_trend == "稳定":
            ml_desc = f"机器学习策略保持在 {int(final_ml)} 元左右，收益率 {ml_return:.2%}。"
        else:
            ml_desc = f"机器学习策略在中期{ml_trend}，最终达到约 {int(final_ml)} 元，收益率 {ml_return:.2%}。"

        description = f"{stock}：\n{trad_desc} {ml_desc}"
        descriptions[stock] = description

    return descriptions


def main():
    token = '1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c'
    pro = initialize_tushare(token)

    # 静态选股：覆盖多个板块，流动性与波动兼顾
    stocks = [
        '600036.SH',  # 招商银行
        '601318.SH',  # 中国平安
        '600519.SH',  # 贵州茅台
        '000858.SZ',  # 五粮液
        '000333.SZ',  # 美的集团
        '300750.SZ',  # 宁德时代
        '603288.SH',  # 海天味业
        '600276.SH'   # 恒瑞医药
    ]
    print("本次回测选股：", stocks)

    # 获取并清洗数据
    start_date, end_date = '20220521', '20250521'
    data_dict = get_stock_data(pro, stocks, start_date, end_date)
    for stock, df in data_dict.items():
        print(f"{stock} 数据范围: {df.index.min().date()} 至 {df.index.max().date()}")

    # 数据统计分析（图 2.1）
    plot_data_analysis(data_dict)

    # 初始化结果存储
    results = {}
    portfolio_values = {}

    # 回测主循环
    for stock, df in data_dict.items():
        # 1. 传统策略
        df = traditional_strategy(df)

        # 2. 机器学习策略
        df, model = ml_strategy_with_tuning(df)

        # 3. 信号示意图（以第一只股票为例）
        if stock == stocks[0]:
            plot_traditional_signals(df, stock)
            plot_ml_signals(df, stock)

        # 4. 回测计算
        trad_ret, trad_sharpe, trad_port = backtest(df, 'traditional')
        ml_ret, ml_sharpe, ml_port         = backtest(df, 'ml')

        # 5. 存储结果
        results[stock] = {
            'Traditional': {'return': trad_ret, 'sharpe': trad_sharpe},
            'ML':          {'return': ml_ret,   'sharpe': ml_sharpe}
        }
        portfolio_values[stock] = {
            'Traditional': trad_port,
            'ML':          ml_port
        }

        print(f"{stock} - 传统策略: 收益率 {trad_ret:.2%}, 夏普 {trad_sharpe:.2f}")
        print(f"{stock} - 机器学习策略: 收益率 {ml_ret:.2%}, 夏普 {ml_sharpe:.2f}")

    # 投资组合价值对比（图 5.1–5.8）
    plot_portfolio_comparison(portfolio_values, data_dict)

    # 特定股票收益率对比（图 4.1）
    plot_specific_returns_comparison(results, stocks[0])

    # 全部收益率和夏普比率对比（图 5.9 & 5.10）
    plot_returns_comparison(results)
    plot_sharpe_comparison(results)

    # 文本报告
    descriptions = generate_portfolio_description(portfolio_values, results)
    for desc in descriptions.values():
        print("\n" + desc)

    # 整体正收益判断
    all_positive = all(
        stat['return'] > 0
        for stock_stats in results.values()
        for stat in stock_stats.values()
    )
    print("所有策略均实现正收益。" if all_positive else "部分策略未实现正收益。")


if __name__ == "__main__":
    main()